import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import TensorBoardLogger
import os
import numpy as np
import matplotlib.pyplot as plt
import yaml
import json
from models.fault_diagnosis import DiagnosisModel
from data.bjut import single_domain_dataloader



def train_classifier():
    # 1. 加载配置
    with open('config/train_classifier.yaml', 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    
    config = config['train_classifier']
    
    # 2. 创建实验名称
    exp_name = (f"classifier_{config['domain_name']}_"
               f"lr{config['lr']}_"
               f"bs{config['batch_size']}_"
               f"ep{config['num_epochs']}")
    
    # 3. 创建保存目录
    save_dir = os.path.join(config['save_dir'], exp_name)
    os.makedirs(save_dir, exist_ok=True)
    
    # 保存配置
    config_save_path = os.path.join(save_dir, 'config.json')
    with open(config_save_path, 'w', encoding='utf-8') as f:
        json.dump(config, f, indent=4, ensure_ascii=False)
    print(f"配置已保存到: {config_save_path}")
    
    # 4. 准备数据加载器
    train_loader = single_domain_dataloader(
        data_path=config['data_path'],
        domain_name=config['domain_name'],
        batch_size=config['batch_size'],
        split='train'
    )
    
    val_loader = single_domain_dataloader(
        data_path=config['data_path'],
        domain_name=config['domain_name'],
        batch_size=config['batch_size'],
        split='val'
    )
    
    # 5. 初始化模型
    model = DiagnosisModel(
        input_dim=config['input_dim'],
        num_classes=config['num_classes'],
        lr=float(config['lr'])
    )
    
    
    # 6. 设置回调和日志
    checkpoint_callback = ModelCheckpoint(
        dirpath=save_dir,
        filename='classifier-{epoch:02d}-{val_loss:.4f}',
        save_top_k=3,
        monitor='val_loss',
        mode='min'
    )
    
    logger = TensorBoardLogger(
        save_dir='logs',
        name=exp_name,
        version=None
    )
    
    # 7. 初始化训练器
    trainer = pl.Trainer(
        max_epochs=config['num_epochs'],
        accelerator='gpu',
        devices=[0],
        callbacks=[checkpoint_callback],
        logger=logger,
        log_every_n_steps=10,
        precision=16,
        val_check_interval=0.5
    )
    
    # 8. 开始训练
    trainer.fit(model, train_loader, val_loader)
    
    # 9. 保存最终模型
    final_model_path = os.path.join(save_dir, 'final_model.pt')
    torch.save(model.state_dict(), final_model_path)
    print(f"最终模型已保存到: {final_model_path}")

def evaluate_classifier(checkpoint_path):
    """评估分类器性能"""
    # 1. 加载配置和模型
    with open('config/train_classifier.yaml', 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)['train_classifier']
    
    model = DiagnosisModel(
        input_dim=config['input_dim'],
        num_classes=config['num_classes']
    )
    model.load_state_dict(torch.load(checkpoint_path))
    model = model.cuda()
    model.eval()
    
    # 2. 准备测试数据
    test_loader = single_domain_dataloader(
        data_path=config['data_path'],
        domain_name=config['domain_name'],
        batch_size=config['batch_size'],
        split='test'
    )
    
    # 3. 评估
    correct = 0
    total = 0
    predictions = []
    true_labels = []
    
    with torch.no_grad():
        for x, y in test_loader:
            x = x.cuda()
            y = y.cuda()
            
            logits, _ = model(x)
            preds = torch.argmax(logits, dim=1)
            
            correct += (preds == y).sum().item()
            total += y.size(0)
            
            predictions.extend(preds.cpu().numpy())
            true_labels.extend(y.cpu().numpy())
    
    # 4. 计算并打印指标
    accuracy = correct / total
    print(f"测试集准确率: {accuracy:.4f}")
    
    # 5. 保存评估结果
    results = {
        'accuracy': accuracy,
        'predictions': predictions,
        'true_labels': true_labels
    }
    
    results_path = os.path.join(os.path.dirname(checkpoint_path), 'evaluation_results.json')
    with open(results_path, 'w') as f:
        json.dump(results, f)
    print(f"评估结果已保存到: {results_path}")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', type=str, 
                       choices=['train', 'evaluate'],
                       default='train')
    parser.add_argument('--checkpoint', type=str, 
                       help='Checkpoint path for evaluation')
    args = parser.parse_args()
    
    if args.mode == 'train':
        train_classifier()
    elif args.mode == 'evaluate':
        if args.checkpoint is None:
            raise ValueError("请提供检查点路径用于评估")
        evaluate_classifier(args.checkpoint)
